import pandas as pd
import os
from datetime import datetime, timedelta
import plotly.graph_objects as go
import numpy as np  # 需要导入numpy来使用np.nan和np.where


# --- 数据加载和处理函数 (不变) ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从指定的CSV文件中加载时间范围内的数据。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]
    all_dfs = []
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            try:
                df = pd.read_csv(file_path, usecols=['open_time', 'open', 'high', 'low', 'close'])
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    """简化版的数据处理和筛选函数。"""
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)

    start_ts = pd.to_datetime(start_time_str)
    end_ts = pd.to_datetime(end_time_str)
    df_filtered = df.loc[start_ts:end_ts]
    print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    return df_filtered


# --- 唐安奇通道计算函数 (不变) ---
def calculate_donchian_channels(df: pd.DataFrame, period: int) -> pd.DataFrame:
    """
    计算唐安奇通道。
    使用 .shift(1) 来确保计算只使用过去的数据，避免未来函数。
    """
    if df.empty:
        return df

    df['donchian_upper'] = df['high'].rolling(window=period).max().shift(1)
    df['donchian_lower'] = df['low'].rolling(window=period).min().shift(1)
    df['donchian_middle'] = (df['donchian_upper'] + df['donchian_lower']) / 2

    print(f"已计算周期为 {period} 的唐安奇通道 (已修正，无未来函数)。")
    return df


# --- 新增：寻找突破点的函数 ---
def find_breakout_points(df: pd.DataFrame) -> pd.DataFrame:
    """
    识别收盘价突破唐安奇通道的点，并添加用于绘图的列。
    """
    # 条件1：收盘价高于上轨
    condition_up = df['close'] > df['donchian_upper']
    # 条件2：收盘价低于下轨
    condition_down = df['close'] < df['donchian_lower']

    # 使用 np.where 高效地创建新列
    # 如果满足条件，则记录标记的位置（略高于/低于K线），否则为 NaN
    # NaN 值在绘图时会被自动忽略
    df['breakout_up_marker'] = np.where(condition_up, df['high'] * 1.001, np.nan)
    df['breakout_down_marker'] = np.where(condition_down, df['low'] * 0.999, np.nan)

    print(f"已识别出 {int(condition_up.sum())} 个向上突破点和 {int(condition_down.sum())} 个向下突破点。")
    return df


# --- 绘图函数 (已更新以标注突破点) ---
def plot_donchian_chart(df: pd.DataFrame, chart_title: str):
    """
    绘制K线图、唐安奇通道，并标注收盘价突破点。
    """
    if df.empty:
        print("数据为空，无法绘制图表。")
        return

    fig = go.Figure()

    # 1. 绘制唐安奇通道 (先绘制，作为背景)
    fig.add_trace(go.Scatter(x=df.index, y=df['donchian_upper'], name='Donchian Upper',
                             line=dict(color='orange', width=1)))

    fig.add_trace(go.Scatter(x=df.index, y=df['donchian_lower'], name='Donchian Lower',
                             line=dict(color='orange', width=1),
                             fill='tonexty',
                             fillcolor='rgba(255,165,0,0.1)'))

    fig.add_trace(go.Scatter(x=df.index, y=df['donchian_middle'], name='Donchian Middle',
                             line=dict(color='orange', width=1, dash='dash')))

    # 2. 绘制K线
    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'],
                                 low=df['low'], close=df['close'],
                                 name='OHLC'))

    # --- 新增：绘制突破点标记 ---
    # 向上突破标记 (绿色上三角)
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['breakout_up_marker'],
        name='Close Breakout Up',
        mode='markers',
        marker=dict(symbol='triangle-up', color='green', size=10)
    ))

    # 向下突破标记 (红色下三角)
    fig.add_trace(go.Scatter(
        x=df.index,
        y=df['breakout_down_marker'],
        name='Close Breakout Down',
        mode='markers',
        marker=dict(symbol='triangle-down', color='red', size=10)
    ))

    # 更新布局
    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    period = df.attrs.get('donchian_period', 'N/A')
    full_title = f'{chart_title} ({start_str} to {end_str})<br>Donchian Channels (Period: {period}) with Breakouts'

    fig.update_layout(
        title_text=full_title,
        xaxis_rangeslider_visible=False,
        yaxis_title='Price',
        hovermode='x unified',
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "BTCUSDT"
    TIMEFRAME = "1D"
    START_TIME = "2024-08-01 00:00:00"
    END_TIME = "2025-08-03 00:00:00"

    DONCHIAN_PERIOD = 10
    # ====================================================

    print("--- 唐安奇通道图表程序启动 ---")

    print(f"\n--- 正在加载图表数据 ({START_TIME} to {END_TIME}) ---")
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_filter_data(raw_data_df, START_TIME, END_TIME)

    if not final_df.empty:
        # 1. 计算唐安奇通道
        final_df = calculate_donchian_channels(final_df, period=DONCHIAN_PERIOD)
        final_df.attrs['donchian_period'] = DONCHIAN_PERIOD

        # 2. 寻找突破点
        final_df = find_breakout_points(final_df)

    # 3. 生成图表
    if not final_df.empty:
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_donchian_chart(final_df, chart_title=chart_title_str)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何图表数据，程序退出。请检查您的配置和文件路径。")